#!/usr/bin/env python36
# -*- coding: utf-8 -*-
"""
Created on July, 2018

@author: Tangrizzly
"""

import argparse
import time
import csv
import pickle
import operator
import datetime
import os

# 参数编译器。便于使用命令行传递参数，执行文件
import numpy

parser = argparse.ArgumentParser()
# 添加dataset参数，默认是sample
parser.add_argument('--dataset', default='sample', help='dataset name: diginetica/yoochoose/sample')
# 获取参数
opt = parser.parse_args()
# 打印所有参数
print(opt)

# 默认数据集文件
dataset = 'sample_train-item-views.csv'
# 判断传入的 --dataset 参数值，然后确定要读取的数据集名称
if opt.dataset == 'diginetica':
    dataset = 'train-item-views.csv'
elif opt.dataset == 'yoochoose':
    dataset = 'yoochoose-clicks.dat'

# ====== 开始读取数据 =====
# 此部分就是从文件中读取会话数据，每个会话对应多个项目，且项目按照时间顺序排序
print("-- Starting @ %ss" % datetime.datetime.now())
with open(dataset, "r") as f:
    # 这里对 yoochoose 的列名读取有问题，yoochoose中yoochoose-clicks.dat文件并无列名
    if opt.dataset == 'yoochoose':
        """
            csv.reader(csvfile, dialect='excel', **fmtparams)，返回一个 reader 对象，该对象将逐行遍历 csvfile，
            DictReader(f, fieldnames=None, restkey=None, restval=None, dialect='excel', *args, **kwds)，创建一个对象，
            该对象在操作上类似于常规reader，但是将每行中的信息映射到一个dict，该 dict 的键由 fieldnames 可选参数给出
        """
        reader = csv.DictReader(f, delimiter=',')
    else:
        reader = csv.DictReader(f, delimiter=';')

    # 会话点击序列，字典形式
    sess_clicks = {}
    # 会话日期序列，记录当前会话中的最后一个时间戳
    sess_date = {}
    ctr = 0
    # 当前会话Id
    curid = -1
    # 当前会话日期
    curdate = None
    # 逐行遍历reader，此时的reader每一个行是一个字典，key是文件列头，value是文件中的列值
    for data in reader:
        # 当前行的会话Id，这里原始文件列头为sessionId，而不是session_id，session_id是用户简化之后的一个数据集中的会话id列头
        # ===== sample数据集 =====
        sessid = data['session_id']
        # ===== sample数据集 =====
        # ===== diginetica数据集 =====
        # sessid = data['sessionId']
        # ===== diginetica数据集 =====
        # 设置当前会话中最后一个项目对应的时间戳，curid为上一个会话Id。sessid为当前行id时，才会设置date
        if curdate and not curid == sessid:
            date = ''
            if opt.dataset == 'yoochoose':
                # strptim e格式化 curdat 为后续字符串形式，mktime则将格式化之后的时间转化为自纪元以来的秒数
                date = time.mktime(time.strptime(curdate[:19], '%Y-%m-%dT%H:%M:%S'))
            else:
                date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))
            # 当前会话最后一个项目的日期设置为date
            sess_date[curid] = date
        # 更新当前会话Id
        curid = sessid
        # 读取当前行的项目
        if opt.dataset == 'yoochoose':
            item = data['item_id']
        else:
            # 这里将item视为一个二元组，由 (item_id,timeframe) 组成
            # ===== sample数据集 =====
            item = data['item_id'], int(data['timeframe'])
            # ===== sample数据集 =====
            # ===== diginetica数据集 =====
            # item = data['itemId'], int(data['timeframe'])
            # ===== diginetica数据集 =====
        # 读取当前日期
        curdate = ''
        if opt.dataset == 'yoochoose':
            curdate = data['timestamp']
        else:
            curdate = data['eventdate']

        # 判断当前会话Id在不在已经读取的点击序列中，如果在，则将当前项目
        if sessid in sess_clicks:
            sess_clicks[sessid] += [item]
        else:
            sess_clicks[sessid] = [item]
        # 点击序列的个数
        ctr += 1

    # print('sess_clicks内容。。。。', sess_clicks)
    # sess_clicks内容。。。。{sessionId:[(item_id,timeframe),(item_id,timeframe),....,],...,}
    # {'1': [('81766', 526309), ('31331', 1031018), ('32118', 243569), ('9654', 75848), ('32627', 1112408), ('33043', 173912), ('12352', 329870), ('35077', 390072), ('36118', 487369), ('129055', 991416)],
    # '2': [('3147', 883039), ('32754', 359839), ('100747', 38317), ('32971', 182759), ('10657', 387340), ('35606', 756335), ('35606', 53684), ('36246', 332000), ('36246', 722492), ('196110', 1045018)], '
    print('sess_clicks长度：', len(sess_clicks))

    date = ''
    if opt.dataset == 'yoochoose':
        date = time.mktime(time.strptime(curdate[:19], '%Y-%m-%dT%H:%M:%S'))
    else:
        date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))
        # list(sess_clicks)将字典形式的sess_clicks转化为列表，默认直降sess_clicks的key转化为列表
        for i in list(sess_clicks):
            # operator模块提供的itemgetter函数用于获取对象的哪些维的数据，这里就是对sess_clicks[i]排序，且以第1维数据排序，即以timeframe从小到大排序
            sorted_clicks = sorted(sess_clicks[i], key=operator.itemgetter(1))
            # 排序之后的项目序列，此时不再需要用到timeframe
            sess_clicks[i] = [c[0] for c in sorted_clicks]
    # 设置最后一个会话Id对应的最后一个项目的日期作为该会话对应的日期
    sess_date[curid] = date
print("-- Reading diginetica-data @ %ss" % datetime.datetime.now())
# 上面一部分从开始处理到这个位置，就是从文件中读取会话数据，每个会话对应多个项目，且项目已经按照时间顺序排序
# ====== 结束读取数据 =====

# Filter out length 1 sessions，删除会话长度小于2的序列 =========
for s in list(sess_clicks):
    if len(sess_clicks[s]) == 1:
        del sess_clicks[s]
        del sess_date[s]
# ======== 删除会话长度小于1的序列 =========

# Count number of times each item appears，统计项目的出现次数 =====
iid_counts = {}
for s in sess_clicks:
    seq = sess_clicks[s]
    for iid in seq:
        if iid in iid_counts:
            iid_counts[iid] += 1
        else:
            iid_counts[iid] = 1
# ====== 统计项目的出现次数 =====

# 按照项目出现的次数对项目排序
sorted_counts = sorted(iid_counts.items(), key=operator.itemgetter(1))

# 过滤前会话长度
length = len(sess_clicks)
# ====== 过滤项目，过滤会话 ======
for s in list(sess_clicks):
    curseq = sess_clicks[s]
    # filter() 函数用于过滤序列，过滤掉不符合条件的元素，返回由符合条件元素组成的新列表。
    # 该接收两个参数，第一个为函数，第二个为序列，序列的每个元素作为参数传递给函数进行判断，然后返回 True 或 False，最后将返回 True 的元素放到新列表中。
    # 过滤掉 curseq 中项目出现次数少于5次的项目
    filseq = list(filter(lambda i: iid_counts[i] >= 5, curseq))
    # 如果过滤掉项目之后，会话长度小于2，则删除该会话
    if len(filseq) < 2:
        del sess_clicks[s]
        del sess_date[s]
    else:
        sess_clicks[s] = filseq
print('删除 %d/%d 个项目' % (len(sess_clicks), length))
# ====== 过滤项目，过滤会话 ======

# Split out test set based on dates，开始划分测试集 =======
dates = list(sess_date.items())
maxdate = dates[0][1]

# ====== 查找最大的会话日期 ======
for _, date in dates:
    if maxdate < date:
        maxdate = date
# ====== 查找最大的会话日期 ======

# 7 days for test，Diginetica数据集取最近7天的数据作为测试集，yoochoose数据集取最近一天的数据作为测试集
splitdate = 0
if opt.dataset == 'yoochoose':
    splitdate = maxdate - 86400 * 1  # the number of seconds for a day：86400
else:
    splitdate = maxdate - 86400 * 7

print('Splitting date', splitdate)  # Yoochoose: ('Split date', 1411930799.0)
# 过滤小于splitdate的数据作为训练集
tra_sess = filter(lambda x: x[1] < splitdate, dates)
# 过滤大于splitdate的数据作为测试集
tes_sess = filter(lambda x: x[1] > splitdate, dates)
print('训练集内容。。。。。')
print(tra_sess)
# Split out test set based on dates，结束划分测试集 =======

# Sort sessions by date
# 按照时间戳排序对训练集和测试集排序
tra_sess = sorted(tra_sess, key=operator.itemgetter(1))  # [(session_id, timestamp), (), ]
tes_sess = sorted(tes_sess, key=operator.itemgetter(1))  # [(session_id, timestamp), (), ]
print('训练集长度。。。。', len(tra_sess))  # 186670    # 7966257
print('测试集长度。。。。', len(tes_sess))  # 15979     # 15324
print('训练集前三个元素', tra_sess[:3])
print('测试集前三个元素', tes_sess[:3])
print("-- Splitting train set and test set @ %ss" % datetime.datetime.now())

# Choosing item count >=5 gives approximately the same number of items as reported in paper
# item原始ID映射到1~n
item_dict = {}


# Convert training sessions to sequences and renumber items to start from 1，将训练会话转换为序列并将项目重新编号以从 1 开始
def obtian_tra():
    # 训练集的会话Ids
    train_ids = []
    # 训练的会话序列
    train_seqs = []
    # 训练日期
    train_dates = []
    item_ctr = 1
    # [(session_id, timestamp)]
    for s, date in tra_sess:
        # 会话ID s 对应的会话序列
        seq = sess_clicks[s]
        # 映射之后的会话序列的项目集合
        outseq = []
        # 遍历原始会话序列，如果原始会话序列中的项目已经在item_dic映射中存在，则添加到outseq中；
        # 如果不在item_dic中，则更新item_dic之后再添加到outseq中
        for i in seq:
            if i in item_dict:
                outseq += [item_dict[i]]
            else:
                outseq += [item_ctr]
                # 更新映射
                item_dict[i] = item_ctr
                item_ctr += 1
        # 如果映射之后的序列序列长度小于2，则不添加到训练集中
        if len(outseq) < 2:  # Doesn't occur
            continue
        # 训练集的会话id
        train_ids += [s]
        # 会话序列对应的时间戳
        train_dates += [date]
        # 会话序列的项目集合
        train_seqs += [outseq]
    print('项目个数', item_ctr)  # 43098, 37484
    # 返回训练集中的会话ID集合，会话日期集合（一个会话ID一个），会话序列集合（项目集）
    return train_ids, train_dates, train_seqs


# Convert test sessions to sequences, ignoring items that do not appear in training set，将测试会话转换为序列，忽略未出现在训练集中的项目
# 与 obtian_tra 函数相同的功能，不过这里是应用在测试集上
def obtian_tes():
    test_ids = []
    test_seqs = []
    test_dates = []
    for s, date in tes_sess:
        seq = sess_clicks[s]
        outseq = []
        for i in seq:
            if i in item_dict:
                outseq += [item_dict[i]]
        if len(outseq) < 2:
            continue
        test_ids += [s]
        test_dates += [date]
        test_seqs += [outseq]
    return test_ids, test_dates, test_seqs


# 分别获得训练集数据和测试集数据
tra_ids, tra_dates, tra_seqs = obtian_tra()
tes_ids, tes_dates, tes_seqs = obtian_tes()


def process_seqs(iseqs, idates):
    """
    增强数据
    原始序列S={v1,v2,v3,...,vn}，则前n-1个数据为输入，vn则为目标值。这里进行数据增强，在前面的基础上，增加训练数据，{v1,v2,..,vn-2}为输入，vn-1为目标值。
    以此类推，{v1,v2,...,vn-3}，则vn-2为预测的目标值
    :param iseqs: 过滤后的数据集合
    :param idates: 过滤后的日期集合
    :return:
    out_seqs: 数据增强之后的序列集合，每个序列包含多个项目（例如{v1,v2,v3,...,vn}等，不包括预测目标值
    out_dates: out_seqs 中每个序列对应的日期
    labs: out_seqs 中序列对应的预测目标值
    ids: out_seqs 中序列对应的id集合，并且是原始会话id映射为0~n之后的id，且out_seqs中多个序列对应的id是相同的，因为是相同的会话增强的得到的。
    """
    # 输出序列（项目集合）
    out_seqs = []
    # 输出日期
    out_dates = []
    # 标签集合，
    labs = []
    ids = []
    for id, seq, date in zip(range(len(iseqs)), iseqs, idates):
        # 对seq序列进行增强，添加多个序列和标签
        for i in range(1, len(seq)):
            # 目标项
            tar = seq[-i]
            # 将目标项添加到标签集合中
            labs += [tar]
            # 添加一个序列到结果序列中
            out_seqs += [seq[:-i]]
            out_dates += [date]
            ids += [id]
    # 返回增强之后的结果
    return out_seqs, out_dates, labs, ids


# 对训练集和测试集进行数据增强，开始 ======
tr_seqs, tr_dates, tr_labs, tr_ids = process_seqs(tra_seqs, tra_dates)
te_seqs, te_dates, te_labs, te_ids = process_seqs(tes_seqs, tes_dates)
# 对训练集和测试集进行数据增强，结束 ======
# 分别将训练集和测试集的序列和目标项标签对应起来
tra = (tr_seqs, tr_labs)
tes = (te_seqs, te_labs)
print('数据增强之后的训练集长度.....', len(tr_seqs))
print('数据增强之后的测试集长度.....', len(te_seqs))
# 输出前三行数据
print('训练集和测试集')
print(tr_seqs[:3], tr_dates[:3], tr_labs[:3])
print(te_seqs[:3], te_dates[:3], te_labs[:3])
# 所有会话加一起的总长度
all = 0

for seq in tra_seqs:
    all += len(seq)
for seq in tes_seqs:
    all += len(seq)
# 平均会话长度=会话总长度/会话个数
print('avg length: ', all / (len(tra_seqs) + len(tes_seqs) * 1.0))
if opt.dataset == 'diginetica':
    if not os.path.exists('diginetica'):
        os.makedirs('diginetica')
    """
    Pickle模块常用函数说明：https://blog.csdn.net/coffee_cream/article/details/51754484
    
    pickle.dump(obj, file, [,protocol])，函数的功能：将obj对象序列化存入已经打开的file中。
    参数讲解：
    obj：想要序列化的obj对象。
    file:文件名称。
    protocol：序列化使用的协议。如果该项省略，则默认为0。如果为负值或HIGHEST_PROTOCOL，则使用最高的协议版本。
    
    pickle.load(file)，函数的功能：将file中的对象序列化读出。
    参数讲解：file：文件名称。
    """
    pickle.dump(tra, open('diginetica/train.txt', 'wb'))
    pickle.dump(tes, open('diginetica/test.txt', 'wb'))
    pickle.dump(tra_seqs, open('diginetica/all_train_seq.txt', 'wb'))
elif opt.dataset == 'yoochoose':
    if not os.path.exists('yoochoose1_4'):
        os.makedirs('yoochoose1_4')
    if not os.path.exists('yoochoose1_64'):
        os.makedirs('yoochoose1_64')
    pickle.dump(tes, open('yoochoose1_4/test.txt', 'wb'))
    pickle.dump(tes, open('yoochoose1_64/test.txt', 'wb'))

    split4, split64 = int(len(tr_seqs) / 4), int(len(tr_seqs) / 64)
    print(len(tr_seqs[-split4:]))
    print(len(tr_seqs[-split64:]))

    # 划分 yoochoose1_4，yoochoose1_64
    tra4, tra64 = (tr_seqs[-split4:], tr_labs[-split4:]), (tr_seqs[-split64:], tr_labs[-split64:])
    seq4, seq64 = tra_seqs[tr_ids[-split4]:], tra_seqs[tr_ids[-split64]:]

    pickle.dump(tra4, open('yoochoose1_4/train.txt', 'wb'))
    pickle.dump(seq4, open('yoochoose1_4/all_train_seq.txt', 'wb'))

    pickle.dump(tra64, open('yoochoose1_64/train.txt', 'wb'))
    pickle.dump(seq64, open('yoochoose1_64/all_train_seq.txt', 'wb'))

else:
    if not os.path.exists('sample'):
        os.makedirs('sample')
    pickle.dump(tra, open('sample/train.txt', 'wb'))
    pickle.dump(tes, open('sample/test.txt', 'wb'))
    pickle.dump(tra_seqs, open('sample/all_train_seq.txt', 'wb'))
    numpy.savetxt('sample/train2.txt', tra)

print('Done.')
